{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e151a651",
   "metadata": {},
   "source": [
    "# 卷积神经网络推理 (1)：Direct 卷积的纯 Python 实现、库函数实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "436d166e",
   "metadata": {},
   "source": [
    "> 公开时间：2022-01-03"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "692bc276",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "该工作是第五届 Ubiquant Challenge 量化新星挑战赛|并行程序设计竞赛的初赛入围相关工作。该工作在【天之孔】队长强宜澄的牵头下完成。\n",
    "\n",
    "这份文档是原题背景的介绍。原题是使用 Winograd 算法高效实现 $3 \\times 3$ 卷积核神经网络的推理过程；我们将在之后的文档再对原题进行讨论。\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cdb7529",
   "metadata": {},
   "source": [
    "在最近 10 年，卷积神经网络 (CNN, Convolutional Neural Network) 在图像识别中取得很大的成功。典型的例子会是 VGG16, ResNet, GoogLeNet。深度卷积网络的模型很大，并且大多数时候在作卷积、池化操作，而较少或只在最后进行多层感知，**卷积网络计算时间相对较长**。\n",
    "\n",
    "我们也知道，神经网络方法的学习需要用到反向传播 (BP, Backward Propagation)；这是非常重要的议题，它推动了自动导数、机器学习框架的实现与推广。但在实际的应用情景中，更多地是**针对已经学习完毕的模型，进行推理计算** (Inference)，即网络的正向过程。这种情形可能是手机输入法、信息软件的文字翻译、自动驾驶等等 (或许金融超额收益也是一种情形，但我还不具备理解量化金融的知识 >.<)。它们都要求在获得输入后，以最迅速的方式通过模型计算并给出结果，而这速度**经常是以毫秒 (ms) 而非秒 (sec) 为单位**。快速的程序实现意味着优质的用户体验、意味着生命的安全、意味着稳定的金钱收益。写一个高效的卷积网络程序至关重要。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "916a7df3",
   "metadata": {},
   "source": [
    ":::{admonition} 拯救生命大挑战！(误\n",
    "\n",
    "您是 Tezla 的算法实现工程师。您需要高效率编写一个卷积神经网络，使得汽车能自动识别到高速路上是否有小孩横穿。您能否拯救一个充满未来与希望 (但现在不太负责) 的生命？\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8495c3c3",
   "metadata": {},
   "source": [
    "我们将会分为三份文档，讨论这个问题。\n",
    "\n",
    "- 第一份文档中，我们会\n",
    "    - 简单介绍卷积网络的定义、公式、基础实现，并辅以图像加深理解；\n",
    "    - 以 Python 为主要编程语言，介绍一些有效的高效率纯 Python 实现方法；\n",
    "    - 指出并简单测试常用的库函数实现 (oneDNN, PyTorch)，以及我们自己编写的 C/C++ 程序的效率；\n",
    "    - 指出纯 Python 实现的不足，并对效率改进的思路做简单说明。\n",
    "- 第二份文档中，我们会\n",
    "    - 介绍 Winograd 算法[^Lavin2016]，并对 $F(6, 3)$ 情形作纯 Python 实现；\n",
    "    - 对 Winograd 算法作简单分析；\n",
    "- 第三份文档中，我们会\n",
    "    - 详细介绍初赛工作中基于 x86-64 (允许 AVX-512) CPU 的 C/C++ 优化思路；\n",
    "    - 对该程序的算法复杂度与调试技巧作说明；\n",
    "    - 对其它 C/C++ 优化思路做简短说明。\n",
    "\n",
    "这是应用与实现文档，我们不对卷积网络或 Winograd 算法作原理性的叙述与证明，也不讨论卷积网络的有效性。我们只考察卷积核大小为 $3 \\times 3$ 的情形。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93c3d469",
   "metadata": {},
   "source": [
    "[^Lavin2016]: Lavin, A.; Gray, S. Fast Algorithms for Convolutional Neural Networks. In *2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)*; IEEE: Las Vegas, NV, USA, **2016**; pp 4013–4021. doi: [10.1109/CVPR.2016.435](https://doi.org/10.1109/CVPR.2016.435)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e5fef0ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numba as nb\n",
    "import torch\n",
    "import ctypes, itertools, time\n",
    "import pandas as pd\n",
    "\n",
    "torch.set_grad_enabled(False)\n",
    "np.set_printoptions(6, suppress=True, linewidth=150)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d07fb419",
   "metadata": {},
   "source": [
    ":::{note}\n",
    "\n",
    "本文档使用的 CPU 是 Intel Xeon Gold 6154。本文档仅使用 8 核并行。\n",
    "\n",
    "编译环境为 GNU C++ 10.2.0。尽管 Intel C++ (icpc) 也可以编译，但不是很建议。\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b9357cea",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.set_num_threads(8)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bb69265",
   "metadata": {},
   "source": [
    ":::{admonition} 文档执行指南\n",
    "\n",
    "本文档使用 `ctypes` 链接 C 语言的编译库。若读者需要执行该 Jupyter 笔记本，请不要仅下载该 .ipynb 文件——该文档还需要使用编译后的 C/C++ 代码。读者可以移步到 [gitee: ajz34/winograd6x3](https://gitee.com/ajz34/winograd6x3) 下载源代码与 Jupyter 笔记本，并在**支持 AVX-512** 的机器上使用 cmake 编译。\n",
    "\n",
    "作为程序实现效率的参比，我们还实现了 Intel oneDNN 下的卷积网络。因此在 cmake 中，我们还引入了 Intel oneAPI 的路径。需要编译 C/C++ 程序的读者可能需要更改该路径。\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc5abbb7",
   "metadata": {},
   "source": [
    ":::{warning}\n",
    "\n",
    "文档作者当前工作 (计算化学理论的双杂化泛函分支) 与该文档所使用的应用或算法都毫无干系。电子积分、格点积分分支或许会使用到相似的算法思想。文档除了卷积网络本身，大多数知识都是现学的，因此可能会在叙述中存在基本性错误。\n",
    "\n",
    "但我希望能强调算法与架构在程序设计中的重要性，读者受众应可以是行外人 (因为作者也是外行 2333)。该系列文档的算法绝非是最高效；它只是程序思路相对简单，效率尚可。\n",
    "\n",
    "对于 Python 语言而言，由于其语言本身的特性，无法达到令人满意的效率。但即使如此，借用合适的库 (NumPy 或 Numba)、调用合适的函数，可以在一定程度上改进效率；甚至可以逼近 C/C++ 程序。这份文档的 Python 实现或许不那么惊艳，但足够表明合适的库与库函数是多么重要。\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1a55768",
   "metadata": {},
   "source": [
    "## 卷积神经网络简介与实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99ba6054",
   "metadata": {},
   "source": [
    "### 3x3 卷积神经网络定义"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4be7d1b8",
   "metadata": {},
   "source": [
    "我们先考虑一种特定的、较小型的情形。现在我们要处理一张图片，其参数为\n",
    "- 高度 (in height) `IH` $H_\\mathrm{in} = 14$；\n",
    "- 宽度 (in width) `IW` $W_\\mathrm{in} = 20$；\n",
    "- 通道数 (in channel) `IC` $C_\\mathrm{in} = 4$。\n",
    "高度、宽度是我们立即能理解的。通道对于普通的图片而言，可以是红、绿、蓝 (RGB 3 通道)，或青、品红、黄、黑 (CMYK 4 通道)、或黑白 (灰度 1 通道)。但这里所提及的图片是广义上的图片，通道数可以是任意的。\n",
    "\n",
    "对上述图片进行卷积操作时，需要通过所谓的“卷积核”。卷积核的作用相当于提取输入图像的特征。其参数为\n",
    "- 高方向 (kernel height) $K_\\mathrm{H} = 3$、宽方向 (kernel width) $K_\\mathrm{W} = 3$；我们会说这类卷积核的大小是 $3 \\times 3$，且后续文档仅讨论这类卷积核；\n",
    "- 输入通道 (in channel) `IC` $C_\\mathrm{in} = 4$，这必须要与输入图片通道数相同；\n",
    "- 输出通道 (out channel) `OC` $C_\\mathrm{out} = 16$。若熟悉 Photoshop，像图像的色调、亮度、梯度、轮廓等等近邻局域信息都可以是输出；当然，在神经网络中，这些都是将会被学习的隐含量，未必是直观的信息。\n",
    "\n",
    "由此，输出图片的参数信息就是确定的了：\n",
    "- 高度 (out height) `OH` $H_\\mathrm{out} = H_\\mathrm{in} - K_\\mathrm{H} + 1 = H_\\mathrm{in} - 2 = 12$；\n",
    "- 宽度 (out width) `OW` $W_\\mathrm{out} = W_\\mathrm{in} - K_\\mathrm{W} + 1 = W_\\mathrm{in} - 2 = 18$；\n",
    "- 通道数 (out channel) `OC` $C_\\mathrm{out} = 16$ 随着卷积核参数固定。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cde83d7",
   "metadata": {},
   "source": [
    "很多时候，如果一次性处理多张图片，程序效率通常会高一些。我们在这里给出分批数量 $N = 8$。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a824f80",
   "metadata": {},
   "source": [
    "图片的储存方式在不同的程序中可能不同。我们所采用的对图片的存储方式是“NCHW”方式，即\n",
    "- 输入图片 `image` $d_{i,c,x,y}$ 维度 $(N, C_\\mathrm{in}, H_\\mathrm{in}, W_\\mathrm{in})$；\n",
    "- 输出图片 `result` $Y_{i,k,x,y}$ 维度 $(N, C_\\mathrm{out}, H_\\mathrm{out}, W_\\mathrm{out})$。\n",
    "\n",
    "卷积核 `filtr` $g_{k,c,u,v}$ 的维度在各种程序通常一致，即 $(C_\\mathrm{out}, C_\\mathrm{in}, K_\\mathrm{H}, K_\\mathrm{W})$ 即 $(C_\\mathrm{out}, C_\\mathrm{in}, 3, 3)$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e9078f5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 8\n",
    "IH, IW = 14, 20\n",
    "IC, OC = 4, 16\n",
    "OH, OW = IH - 2, IW - 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "40445217",
   "metadata": {},
   "outputs": [],
   "source": [
    "image  = np.random.random(( N, IC, IH, IW)).astype('f') * 255\n",
    "filtr  = np.random.random((OC, IC,  3,  3)).astype('f')\n",
    "result = np.zeros        (( N, OC, OH, OW)).astype('f')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c266a2c",
   "metadata": {},
   "source": [
    "为了简化后续的程序，我们会将输入图像与卷积核放在下述被折叠的代码中，以类 `CNNInput` 表示，并且给出各种维度信息。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "eba9edec",
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "class CNNInput:\n",
    "    \n",
    "    def __init__(self, image, filtr):\n",
    "        self.image = image\n",
    "        self.filtr = filtr\n",
    "        # Dimensions that could be infered from input image and filter\n",
    "        self.N, self.IC, self.IH, self.IW = image.shape\n",
    "        self.OC = filtr.shape[0]\n",
    "        self.OH, self.OW = self.IH - 2, self.IW - 2\n",
    "        # Check sanity for filter dimension\n",
    "        assert(len(filtr.shape) == 4)\n",
    "        assert(self.IC == filtr.shape[1])\n",
    "        assert(filtr.shape[2] == filtr.shape[3])\n",
    "        \n",
    "    @property\n",
    "    def dim_result(self):\n",
    "        \"\"\"Return expected dimension of output image (as result)\"\"\"\n",
    "        return self.N, self.OC, self.OH, self.OW"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b53e2510",
   "metadata": {},
   "source": [
    "### 卷积网络过程图示、公式与极简 Python 实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f29d62d3",
   "metadata": {},
   "source": [
    "卷积的过程相当于对每个输出通道 $k$，将图像分块 $\\boldsymbol{d}$ 与卷积核 $\\boldsymbol{g}$ 作乘积；随后输出一个像素变小了一些的图像 $\\boldsymbol{Y}$。若卷积核大小是 $3 \\times 3$，则输出图像相对于输入图像，在宽与高上都各少 2 个像素。\n",
    "\n",
    "下述图片是基于分批数量 $N = 1$ 绘制；不过我们的程序实现中 $N = 8$。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "215e6690",
   "metadata": {},
   "source": [
    "![cnn_direct](figures/cnn_direct.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cfcbca3",
   "metadata": {},
   "source": [
    "若从公式上来理解，则可以写为\n",
    "\n",
    "$$\n",
    "Y_{i,k,x,y} = \\sum_{c}^{C_\\mathrm{in}} \\sum_{u}^{K_\\mathrm{H}} \\sum_{v}^{K_\\mathrm{W}} d_{i,c,x+u,y+v} g_{k,c,u,v}\n",
    "$$\n",
    "\n",
    "上式中，非求和角标有分批数量 $i < N$、输出通道 $k < OC$、输入图像像素数 $x < OH$ 与 $y < OW$；被求和角标有输入通道 $c < IC$、以及卷积核 $u < KH = 3$ 与 $v < KW = 3$。依下述使用 NumPy 程序 `conv_direct_python_7loops`，利用对所有角标 $i, k, x, y, c, u, v$ 的**七重循环**，就可以实现卷积过程。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e3a105aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_direct_python_7loops(image, filtr):\n",
    "    inp = CNNInput(image, filtr)\n",
    "    result = np.zeros(inp.dim_result).astype('f')\n",
    "    for i in range(inp.N):               # batch size\n",
    "        for k in range(inp.OC):          # out channel\n",
    "            for x in range(inp.OH):      # out height\n",
    "                for y in range(inp.OW):  # out width\n",
    "                    for c in range(inp.IC):     # input channel\n",
    "                        for u in range(3):      # kernel width\n",
    "                            for v in range(3):  # kernel height\n",
    "                                result[i, k, x, y] += image[i, c, x+u, y+v] * filtr[k, c, u, v]\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2fd8d2da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "654 ms ± 1.39 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 7 -n 1\n",
    "result = conv_direct_python_7loops(image, filtr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "981c7418",
   "metadata": {},
   "source": [
    "但 Python 是一种很神奇的语言：它是以 C 作为底层实现的语言。若**库函数使用得当**，一般来说还是**会有比较可观的效率**，并且**大大简化代码量**。首先，注意到 Python 的 for 循环效率是恶名昭著的；它可以通过 `itertools` 进行改进。在我们测试的机器上，一般有 40 ms 的提升 (但也只有 40 ms)。更关键的是，使用 `itertools` 可以减少 Python 的强制缩进数，使代码看起来更加友好，行数更少。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cec70c04",
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_direct_python_7loops(image, filtr):\n",
    "    inp = CNNInput(image, filtr)\n",
    "    result = np.zeros(inp.dim_result).astype('f')\n",
    "    for i, k, x, y, c, u, v in itertools.product(\n",
    "            range(inp.N), range(inp.OC), range(inp.OH), range(inp.OW),\n",
    "            range(inp.IC), range(3), range(3)):\n",
    "        result[i, k, x, y] += image[i, c, x+u, y+v] * filtr[k, c, u, v]\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "03625599",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "611 ms ± 582 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 7 -n 1\n",
    "result_ref = conv_direct_python_7loops(image, filtr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "493121b0",
   "metadata": {},
   "source": [
    "看起来这段代码毫无破绽，**没有多余的循环、没有多余的浮点计算**。**真的是这样吗？是也不是**。原理上，这句话是对的；但我们马上就会发现，上面这段代码的效率是有多糟糕。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bc573ec",
   "metadata": {},
   "source": [
    "## 纯 Python 性能优化"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76bd8167",
   "metadata": {},
   "source": [
    "### 利用 NumPy 减少代码循环数"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6ff7a88",
   "metadata": {},
   "source": [
    "我们知道 NumPy 是 Python 中专门对矩阵或张量进行优化的库。如果合理利用 NumPy，可以大大加速程序效率，同时降低代码复杂程度。\n",
    "\n",
    "但图像卷积问题，麻烦就麻烦它与矩阵计算神似，却并不能直接化为矩阵或张量计算问题。但如果我们不考虑输出像素角标 $x, y$ (我们把这两个角标放到右上方)，那么它就化为了简单的张量数乘与求和运算：\n",
    "\n",
    "$$\n",
    "Y_{i,k}^{(x,y)} = \\sum_{c}^{C_\\mathrm{in}} \\sum_{u}^{K_\\mathrm{H}} \\sum_{v}^{K_\\mathrm{W}} d^{(x,y)}_{i,c,u,v} g_{k,c,u,v}\n",
    "$$\n",
    "\n",
    "我们可以对 $i, k, x, y$ 进行显式循环；而剩下的乘积与求和就交给 NumPy 完成。需要指出，**整个程序的核心代码仅两行** (Python/NumPy 是真的方便 =ω=)。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "964b1c56",
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_direct_numpy_4loops(image, filtr):\n",
    "    inp = CNNInput(image, filtr)\n",
    "    result = np.empty(inp.dim_result).astype('f')\n",
    "    for i, k, x, y in itertools.product(range(inp.N), range(inp.OC), range(inp.OH), range(inp.OW)):\n",
    "        result[i, k, x, y] = (image[i, :, x:x+3, y:y+3] * filtr[k]).sum()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2114f2a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "111 ms ± 288 µs per loop (mean ± std. dev. of 7 runs, 7 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 7 -n 7\n",
    "result = conv_direct_numpy_4loops(image, filtr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44c50cc0",
   "metadata": {},
   "source": [
    "### 利用 einsum 张量缩并提升性能"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17899ebb",
   "metadata": {},
   "source": [
    "我们看到了 **6 倍左右的效率提升**！但上述程序的效率非常低。\n",
    "\n",
    "借助于强大且直观的 NumPy 张量缩并引擎 `np.einsum`，我们可以将 (右下角的) 角标缩并顺序写为字符串。计算过程**仍然可以在两行之内高效并行地解决** (记得要将 `optmize=True` 的选项打开)，并且清晰地展示了张量缩并的具体过程；但仍然需要对 $x, y$ 作循环："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e777af10",
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_direct_numpy_einsum(image, filtr):\n",
    "    inp = CNNInput(image, filtr)\n",
    "    result = np.empty(inp.dim_result).astype('f')\n",
    "    for x, y in itertools.product(range(inp.OH), range(inp.OW)):\n",
    "        result[:, :, x, y] = np.einsum(\"icuv, kcuv -> ik\", image[:, :, x:x+3, y:y+3], filtr, optimize=True)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1fa34d4a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "13.6 ms ± 132 µs per loop (mean ± std. dev. of 7 runs, 7 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 7 -n 7\n",
    "result = conv_direct_numpy_einsum(image, filtr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3d64ada",
   "metadata": {},
   "source": [
    "### 将问题化简为矩阵乘法"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7347699c",
   "metadata": {},
   "source": [
    "我们又看到了 **6 倍左右效率的提升**！尽管 `np.einsum` 通常可以找到一条最佳的优化路径，但其实现可能是通过 `np.tensordot` 而非 `np.matmul` 给出。如果我们注意到卷积可以写成类似于矩阵乘法的形式：\n",
    "\n",
    "$$\n",
    "Y_{i,k}^{(x,y)} = \\sum_{c}^{C_\\mathrm{in}} \\sum_{u}^{K_\\mathrm{H}} \\sum_{v}^{K_\\mathrm{W}} d^{(x,y)}_{i,cuv} g_{k,cuv}\n",
    "$$\n",
    "\n",
    "那么实际上，$c, u, v$ 作为被缩并角标，可以用矩阵乘法处理。这**仍然可以在两行代码中完成**，但代码可能不是很直观："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "50495b79",
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_direct_numpy_matmul(image, filtr):\n",
    "    inp = CNNInput(image, filtr)\n",
    "    result = np.empty(inp.dim_result).astype('f')\n",
    "    for x, y in itertools.product(range(inp.OH), range(inp.OW)):\n",
    "        result[:, :, x, y] = image[:, :, x:x+3, y:y+3].reshape(inp.N, -1) @ filtr.reshape(inp.OC, -1).T\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7295145c",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.04 ms ± 13.7 µs per loop (mean ± std. dev. of 7 runs, 7 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 7 -n 7\n",
    "result = conv_direct_numpy_matmul(image, filtr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1e66b5f",
   "metadata": {},
   "source": [
    "但依这个思路，**还能有更快的代码**。我们注意到卷积核 `filtr` $g_{k,cuv}$ 一直被用到，但在迭代过程中始终需要经过转置并压平成为 $(g^\\dagger)_{cuv,k}$。我们可以**提前分配额外的内存，卷积核转置**到 `trans_filtr`。当前的问题太小，可能无法体现明显的效率优势。**痛心疾首 (sic) 的是，由于转置多了一行代码，因此代码量增加了 50%**。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "442166b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_direct_numpy_matmul(image, filtr):\n",
    "    inp = CNNInput(image, filtr)\n",
    "    result = np.empty(inp.dim_result).astype('f')\n",
    "    trans_filtr = filtr.reshape(inp.OC, -1).T.copy()\n",
    "    for x, y in itertools.product(range(inp.OH), range(inp.OW)):\n",
    "        result[:, :, x, y] = image[:, :, x:x+3, y:y+3].reshape(inp.N, -1) @ trans_filtr\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d240f46c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "930 µs ± 35.5 µs per loop (mean ± std. dev. of 7 runs, 7 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 7 -n 7\n",
    "result = conv_direct_numpy_matmul(image, filtr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "567f2708",
   "metadata": {},
   "source": [
    "我们将程序从一开始的约 650 ms，经过 numpy 数乘与求和的引入到约 130 ms，再是 einsum 的引入到约 15 ms，最终依靠矩阵乘法提升到了 1 ms。**这是约 600 倍的效率提升**！但若只考虑并行的增益，那么效率提升最多也只有 8 倍 (因为我们目前只开了 8 核并行)。我们明确地看到，即使浮点计算数 FLOP 相同，若使用不同的算法或库函数，就会产生完全不同的效果。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "869d1277",
   "metadata": {},
   "source": [
    "### 七重循环的 Numba 实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ab4f9a8",
   "metadata": {},
   "source": [
    "我们最后指出，纯 Python 方式高效实现还有一种途径，即使用 Numba。Numba 通过 LLVM (应用于各种高效能编译器的技术)，实现纯 Python 代码的编译；其执行效率与 C/C++ 代码类似。我们不妨尝试一下，如果抛开 Python 语言本身的包袱，这段代码能有多快？"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b60d8086",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "@nb.jit(\"float32[:,:,:,::1](float32[:,:,:,::1], float32[:,:,:,::1])\", nopython=True, parallel=True)\n",
    "def conv_direct_numba_7loops(image, filtr):\n",
    "    # === preparation ===\n",
    "    N, IC, IH, IW = image.shape\n",
    "    OC = filtr.shape[0]\n",
    "    OH, OW = IH - 2, IW - 2\n",
    "    result = np.zeros(( N, OC, OH, OW), dtype=np.float32)\n",
    "    # === direct algorithm ===\n",
    "    for i in nb.prange(N):           # batch size [parallelized]\n",
    "        for k in range(OC):          # out channel\n",
    "            for x in range(OH):      # out height\n",
    "                for y in range(OW):  # out width\n",
    "                    for c in range(IC):     # input channel\n",
    "                        for u in range(3):      # kernel width\n",
    "                            for v in range(3):  # kernel height\n",
    "                                result[i, k, x, y] += image[i, c, x+u, y+v] * filtr[k, c, u, v]\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7232e277",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "361 µs ± 69.9 µs per loop (mean ± std. dev. of 7 runs, 7 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -r 7 -n 7\n",
    "result_ref = conv_direct_numba_7loops(image, filtr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f76a920a",
   "metadata": {},
   "source": [
    "说来可能不信，Numba 的速度到了 400 µs 级别。它相比于 Python 的七重循环实现，在算法相同的情况下，**效率提升了 1500 倍**！如此大的效率提升，多少有其偶然性。总结来说，Python 难以做到但 Numba 或编译语言可以实现的是\n",
    "\n",
    "1. Numba 实现中，我们对分批数量 $N$ 作简单并行。我们目前开了 8 核并行，因此效率提升也可能有 8 倍左右。\n",
    "2. Numba 的实现相当于 C/C++ 编译语言了。对于 x86-64 架构，许多浮点计算在比较激进的编译指令 (譬如 `-O3`, `-march`) 会自动“向量化”(vectorized) 地进行；譬如 AVX-512 指令集可以一次性对 16 个浮点数作加法或乘法。因此，效率最大可以提升 16 倍；但很多时候能达到 4 倍就已经很令人满意了。\n",
    "3. AVX 中还单独设“乘积累加”(Fused Multiply Add/Accumulate) 运算，即本来要分两次执行的加法与乘法运算放在一次执行。卷积问题完美契合这个运算，因此还能获得 2 倍额外加速。\n",
    "4. AVX (或其子集 SSE) 依靠向量化指令，对循环也能加速。\n",
    "5. 当前问题的输入图像、卷积核、输出图像都可以储存在 L2 cache 中。或许 Numba 对 cache 的处理要比 NumPy 好一些。\n",
    "6. 说不清楚的其它原因都归到 Python 的语言包袱吧 ≥ω≤\n",
    "\n",
    "但同时也要说明，**一旦卷积问题变大**，图像、卷积核无法存入 L2 cache 时，**这个算法会立即失败**。在下面的测评环节就能看出端倪。\n",
    "\n",
    "关于 AVX 指令集的一些操作，我们还会在第二篇文档中作展开。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d1ab260",
   "metadata": {},
   "source": [
    "## 库函数介绍与初步性能测评"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7395218d",
   "metadata": {},
   "source": [
    "### 简化的实际问题：VGG16 conv(3.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60c00094",
   "metadata": {},
   "source": [
    "为了考察较为现实的网络效率，同时介绍各个库函数实现，我们先取 VGG16[^Simonyan2015] 的一层网络 conv(3.2) 作为这一节的讨论对象。该网络还是有一点点挑战性的。其\n",
    "- 图像大小 $H_\\mathrm{in} = W_\\mathrm{in} = 56$；\n",
    "- 输入、输出通道数 $C_\\mathrm{in} = C_\\mathrm{out} = 256$；\n",
    "- 每批图像数量为 $N = 8$。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fbbe867",
   "metadata": {},
   "source": [
    "[^Simonyan2015]: Simonyan, K.; Zisserman, A. Very Deep Convolutional Networks for Large-Scale Image Recognition. *ICLR 2015*. arXiv: [1409.1556](http://arxiv.org/abs/1409.1556) [cs]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6aefea14",
   "metadata": {},
   "outputs": [],
   "source": [
    "image  = np.random.random((  8, 256, 56, 56)).astype('f') * 255\n",
    "filtr  = np.random.random((256, 256,  3,  3)).astype('f')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5b2f90b",
   "metadata": {},
   "source": [
    "### 效率测评量标：GFLOPs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0be8d4d",
   "metadata": {},
   "source": [
    "为了测评程序的效率，我们需要了解卷积过程的运算数量 (FLoating point OPeration, FLOP)，以及当前程序的执行速度 (Giga-FLOP/sec 或 GFLOPs)。\n",
    "\n",
    "计算 FLOP 的方式很简单。首先，这是一个简单张量缩并过程[^note-contraction]。我们看到等式左边作为输出图像的 $Y_{i,k,x,y}$，它有四个角标，其维度分别对应 $N, C_\\mathrm{out}, H_\\mathrm{out}, W_\\mathrm{out}$；而等式右边被求和的角标为 $c, u, v$，分别对应 $C_\\mathrm{in}, K_\\mathrm{H}=3, K_\\mathrm{W}=3$。同时，该式同时需要进行乘法与加法；假若算机的浮点乘法与加法运算效率相同[^note-mult-add-effc]，那么 FLOP 还需要乘以 2 (其 1 是一次加法、其 1 是一次乘法)。但需要注意的是，分批数量 $N$ 与网络参数无关，因此要将其去除。最终的 FLOP 计算式为\n",
    "\n",
    "$$\n",
    "\\mathrm{FLOP} = 2 C_\\mathrm{out} H_\\mathrm{out} W_\\mathrm{out} C_\\mathrm{in} K_\\mathrm{H} K_\\mathrm{W} = 18 C_\\mathrm{out} H_\\mathrm{out} W_\\mathrm{out} C_\\mathrm{in}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ae4c75f5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3439853568"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_FLOP(inp: CNNInput):\n",
    "    return 18 * inp.OC * inp.OH * inp.OW * inp.IC\n",
    "\n",
    "CNNInput.get_FLOP = get_FLOP\n",
    "CNNInput(image, filtr).get_FLOP()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9b24161",
   "metadata": {},
   "source": [
    "即 VGG16 conv(3.2) 的 Giga-FLOP 或 GFLOP 数是 3.44。\n",
    "\n",
    "我们可以使用下面折叠代码的函数 `evaluate_GFLOPs`，给出网络的参数、浮点运算数、以及当前实现下的程序效率。**GFLOP 与执行时间都除去了分批数量 $N$ 的贡献**。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "73656b99",
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def evaluate_GFLOPs(image, filtr, cnn, loops=10, preloops=3, flag_print=True):\n",
    "    inp = CNNInput(image, filtr)\n",
    "    flop = inp.get_FLOP()\n",
    "    if flag_print:\n",
    "        print(\"--- (N,IC,IH,IW,OC): ({:3d}, {:3d}, {:3d}, {:3d}, {:3d}); GFLOP    : {:9.3f}\".format(\n",
    "              inp.N, inp.IC, inp.IH, inp.IW, inp.OC, flop / 1000**3))\n",
    "    tlist = []\n",
    "    for _ in range(preloops):\n",
    "        cnn(image, filtr)\n",
    "    for _ in range(loops):\n",
    "        t0 = time.time(); cnn(image, filtr); tlist.append(time.time() - t0)\n",
    "    trun, tstd = np.sum(tlist), np.std(tlist)\n",
    "    time_ms = trun / loops / inp.N * 1000\n",
    "    time_std = tstd / inp.N * 1000\n",
    "    gflops = flop / time_ms / 1000**2\n",
    "    if flag_print:\n",
    "        print(\"---    ms/run/batch: {:9.3f} ± {:<7.3f}      ; GFLOP/sec: {:9.3f}\".format(time_ms, time_std, gflops))\n",
    "    return flop, time_ms, gflops"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7f48aca",
   "metadata": {},
   "source": [
    "譬如对于 NumPy 的矩阵乘法方案，其执行时间大约在 70 ms 左右，对应程序效率大约是 50 GFLOPs 或 GFLOP/sec。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "810c18db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- (N,IC,IH,IW,OC): (  8, 256,  56,  56, 256); GFLOP    :     3.440\n",
      "---    ms/run/batch:    60.285 ± 1.836        ; GFLOP/sec:    57.060\n"
     ]
    }
   ],
   "source": [
    "evaluate_GFLOPs(image, filtr, conv_direct_numpy_matmul);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b40904e2",
   "metadata": {},
   "source": [
    "[^note-contraction]: 事实上有相对更复杂的张量缩并过程。譬如我们若要考察基 $\\{p,q,\\cdots\\}$ 的矩阵 $F_{pq}$ 在另一组基 $\\{\\mu, \\nu, \\cdots\\}$ 下的转换矩阵 $F_{\\mu \\nu}$ (两组基的转换矩阵是 $C_{\\mu p}$)，即 $F_{\\mu \\nu} = \\sum_{p,q} C_{\\mu p} F_{pq} C_{\\nu q}$ 基的空间大小是 $N$。如果用简单缩并的思路考察，我们知道等式左边有两个角标 $\\mu, \\nu$ 分别对应维度是 $N, N$，等式右边被求和角标 $p, q$ 分别对应维度是 $N, N$。且每一次求和都要经过两次乘法，因此 FLOP 还需要乘以 3。最终 FLOP 式为 $\\mathrm{FLOP(naive)} = 3 N^2 N^2 = 3 N^4$。但如果我们将该转换过程拆分为两步，则就变成两次矩阵乘法了：$G_{\\mu q} = \\sum_p C_{\\mu p} F_{pq}$, $F_{\\mu \\nu} = \\sum_q G_{\\mu q} C_{\\nu q}$；那么这种“复杂”过程的 FLOP 就变为 $\\mathrm{FLOP(complicated)} = 2 N^3 + 2 N^3 = 4 N^3$。因此，并非所有张量缩并过程都适合用正文所述的方法分析；只是 Direct 卷积方法是一种“简单”缩并过程，这样分析没有问题。同时参考 Lavin 2016 文章[^Lavin2016]的 eq(19) 与 Table 3。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4accf63",
   "metadata": {},
   "source": [
    "[^note-mult-add-effc]: 我们或许会听到，若要计算 `b=a*2.0` 的计算，比起使用浮点数乘法，更适合用加法 `b=a+a`。若出于减小计算误差，这确实可能是有效的；但若说加法效率一般高于乘法，则很可能是不正确的——这应是与芯片构造有关。高效计算一般需要用到向量化运算。如果我们看 Intel Intrinsics Guide 的说明，会发现 avx2 级别的 32 位浮点加法 ([\\_mm256_add_ps](https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_add_ps&ig_expand=4956,4956,156,156)) 与乘法 ([\\_mm256_mul_ps](https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_mul_ps&ig_expand=4956))、甚至与乘积累加 ([\\_mm256_fmadd_ps](https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_fmadd_ps&ig_expand=3202)) 在最新的架构下 (Icelake, Skylake) 效能完全相同。不过早一些的架构还是加法比乘法比乘积累加快，只是区别没有那么明显。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05ceaadf",
   "metadata": {},
   "source": [
    "### 真实环境的卷积效率 (1)：Python 实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88058827",
   "metadata": {},
   "source": [
    "我们先看看我们的纯 Python 程序效率。\n",
    "\n",
    "- 即使使用了 NumPy 的四重循环方法，耗时还是太长，要跑 8 sec。VGG16 可是有 16 层啊，要每层都是这效率，那总耗时是 $8 \\times 16 = 128$ sec。用这种卷积神经网络做自动驾驶，那这就好比车开在路上，撞死了一个孩子之后要花两分钟才能反应过来。写出这种代码的人可以去枪毙两分钟。如果 Tezla 公司采用了这种代码，那这个公司不可能倒闭：它连能倒闭的资本都不会有 (手动狗头)。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "f1857a1e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- (N,IC,IH,IW,OC): (  8, 256,  56,  56, 256); GFLOP    :     3.440\n",
      "---    ms/run/batch:  6594.200 ± 0.000        ; GFLOP/sec:     0.522\n"
     ]
    }
   ],
   "source": [
    "evaluate_GFLOPs(image, filtr, conv_direct_numpy_4loops, loops=1, preloops=0);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1d23f52",
   "metadata": {},
   "source": [
    "- 同样是七重循环，若用 Numba 进行编译加速会快许多，但仍然远远不达标："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e316ca05",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- (N,IC,IH,IW,OC): (  8, 256,  56,  56, 256); GFLOP    :     3.440\n",
      "---    ms/run/batch:   234.823 ± 0.250        ; GFLOP/sec:    14.649\n"
     ]
    }
   ],
   "source": [
    "evaluate_GFLOPs(image, filtr, conv_direct_numba_7loops, loops=3, preloops=1);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68cc30d5",
   "metadata": {},
   "source": [
    "如果 Tezla 电动车以 100 km/h (28 m/sec) 的速度开在路上，发现 150 m 前有一小孩横穿马路 (高速路行车距离一般是 100 m、故障标志需放在 150 m 开外)。一般汽车在 100 km/h 下的紧急制动距离是 42 m。如果 Tezla 看到前面小孩，需要 0.3 sec 反应过来，那么就因算法延迟额外地增加了 $0.3 \\times 28 = 8.4$ m 的制动距离。需要指出这才只是 VGG16 网络的一层；如果全部 16 层都是这个效率，那总制动距离就变成 $8.4 \\times 16 + 28 = 162.4 > 150$ m。*拯救生命挑战* 大失败！"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d0e054a",
   "metadata": {},
   "source": [
    "- 利用 `np.einsum` 的效率会有明显提升："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e547372d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- (N,IC,IH,IW,OC): (  8, 256,  56,  56, 256); GFLOP    :     3.440\n",
      "---    ms/run/batch:    82.997 ± 0.588        ; GFLOP/sec:    41.446\n"
     ]
    }
   ],
   "source": [
    "evaluate_GFLOPs(image, filtr, conv_direct_numpy_einsum);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fb2b437",
   "metadata": {},
   "source": [
    "- 使用矩阵乘法的效率会更高："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b3c8a672",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- (N,IC,IH,IW,OC): (  8, 256,  56,  56, 256); GFLOP    :     3.440\n",
      "---    ms/run/batch:    57.045 ± 0.507        ; GFLOP/sec:    60.300\n"
     ]
    }
   ],
   "source": [
    "evaluate_GFLOPs(image, filtr, conv_direct_numpy_matmul);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "709c2204",
   "metadata": {},
   "source": [
    "不过需要指出，即使用 NumPy 程序确实可以拯救小孩的生命，但因卷积网络导致的额外制动距离仍然有约 30 m。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a83a9de1",
   "metadata": {},
   "source": [
    "真的是因为 Python，所以效率低下吗？不全然是。\n",
    "\n",
    "- 首先，简单的纯 Python 实现，效率到这里我想已经是极限了。我相信仍然有可以提升的空间，譬如在复杂算法下，使用 Numba 或许能在手动调整算法后进一步加速。但 Numba 对性能微调的要求较高，因此想要写出真正高效的程序，非常困难。事实上，我在写这份文档时尝试过 Numba 的实现，但并未成功。\n",
    "- 我们以前看到过，在小型问题中，Numba 的实现效率奇高；但为何一到 VGG16 conv(3.2) 就如此糟糕呢？\n",
    "    - 这是因为一旦图像、卷积核变大后，就无法在 L2 或 L3 cache、而要在主内存 (Dynamic Random Access Memory) 中储存这些张量；而 DRAM 的索引效率是低得恐怖的。程序看起来一直在 800% CPU 高速运转，也没有浪费任何浮点运算，但实际上却一直在忙着从主内存索取数据。\n",
    "    - 这是影响程序效率的最关键因素，但并非不可避免。通过缓存友好 (cache aware/friendly) 的编程方式，这种问题是可以避免的。我们会在以后的文档提及这个问题。\n",
    "- 这不意味着 Python 一无所是。事实上，上述 Python 效率低下的原因应该是代码并非完全缓存友好 (但涉及矩阵计算的部分是 NumPy 调用 MKL 库的，这部分效率有保证)。如果随便写一个 C/C++ 程序，那会跟这里的 Numba 实现一样，效率也不见得高。我们马上就会看到。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d455a434",
   "metadata": {},
   "source": [
    "### 真实环境的卷积效率 (2)：最简单的 C/C++ 实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bac047d",
   "metadata": {},
   "source": [
    "这里作为一个 baseline，我们考察下述最简单的 C 程序 (`driver.c` 下的 `naive_conv` 函数)。它通过简单的七重循环实现，并且对于分批数 $N$ 作简单并行 (与我们先前的 Numba 实现相同)。它抛弃了 Python 的所有包袱，并且优化选项几乎是最高的 (`-O3 -march=native`)，或许因此其效率比算法完全一致的 Numba 效率要高。但即使如此，其实现效率远不如上述的 Python/NumPy。可见在这个问题中，缓存友好的编程方式是多么重要。\n",
    "\n",
    "我们用 `ctypes` 调用该函数并考察其效率。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "9b35f4a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "clib_winograd_f6x3 = np.ctypeslib.load_library(\"libwinograd_f6x3.so\", \".\")\n",
    "def conv_direct_c_naive(image, filtr):\n",
    "    inp = CNNInput(image, filtr)\n",
    "    result = np.empty(inp.dim_result).astype('f')\n",
    "    clib_winograd_f6x3.naive_conv(\n",
    "        image.ctypes.data_as(ctypes.c_void_p),\n",
    "        filtr.ctypes.data_as(ctypes.c_void_p),\n",
    "        result.ctypes.data_as(ctypes.c_void_p),\n",
    "        ctypes.c_int(inp.N), ctypes.c_int(inp.IC),\n",
    "        ctypes.c_int(inp.IH), ctypes.c_int(inp.IW), ctypes.c_int(inp.OC))\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "16a499cd",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- (N,IC,IH,IW,OC): (  8, 256,  56,  56, 256); GFLOP    :     3.440\n",
      "---    ms/run/batch:   144.460 ± 0.148        ; GFLOP/sec:    23.812\n"
     ]
    }
   ],
   "source": [
    "evaluate_GFLOPs(image, filtr, conv_direct_c_naive, loops=3, preloops=1);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0303be59",
   "metadata": {},
   "source": [
    "### 真实环境的卷积效率 (3)：库函数与我们的程序"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0ccd72c",
   "metadata": {},
   "source": [
    "我们这里需要经常用到函数签名相同的 C 程序，因此先作下述程序定义。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "264b209f",
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "def conv_c_API(clib):\n",
    "    def inner(image, filtr):\n",
    "        inp = CNNInput(image, filtr)\n",
    "        result = np.empty(inp.dim_result).astype('f')\n",
    "        clib.winconv(\n",
    "            image.ctypes.data_as(ctypes.c_void_p),\n",
    "            ctypes.c_int(inp.IH), ctypes.c_int(inp.IW), ctypes.c_int(inp.IC),\n",
    "            filtr.ctypes.data_as(ctypes.c_void_p),\n",
    "            ctypes.c_int(inp.OC), ctypes.c_int(inp.N),\n",
    "            result.ctypes.data_as(ctypes.c_void_p))\n",
    "        return result\n",
    "    return inner"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f29da2d",
   "metadata": {},
   "source": [
    "- **Intel oneDNN** 是开源库；它需要通过 C/C++ 实现。下面是 Direct 卷积 (直接卷积，即上文提到的原理) 实现。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "769ece12",
   "metadata": {},
   "outputs": [],
   "source": [
    "clib_direct_oneDNN = np.ctypeslib.load_library(\"libdnnl_direct.so\", \".\")\n",
    "conv_direct_oneDNN = conv_c_API(clib_direct_oneDNN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "7b2e0735",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- (N,IC,IH,IW,OC): (  8, 256,  56,  56, 256); GFLOP    :     3.440\n",
      "---    ms/run/batch:     3.724 ± 0.158        ; GFLOP/sec:   923.752\n"
     ]
    }
   ],
   "source": [
    "evaluate_GFLOPs(image, filtr, conv_direct_oneDNN, loops=100, preloops=10);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b648b6a",
   "metadata": {},
   "source": [
    "- **Intel oneDNN** 还提供了 Winograd 卷积实现。若不考虑机器精度损失，Winograd 卷积与 Direct 卷积的实现结果完全相同。但与上面提到的种种优化算法或方法有本质上的不同：**Winograd 卷积是降低了实际运算 FLOP 的算法**。该算法有其适用范围与限制，因此并不一定有明显的速度提升，甚至可能不升反降。Winograd 算法会是下两篇文档的主要议题。这里所用的 GFLOP/sec 是**有效 GFLOP 即 Direct 卷积的浮点运算数**，而非 Winograd 卷积实际的浮点运算数。Intel oneDNN 实现的 Winograd 应是 $F(2, 3)$ 或 $F(4, 3)$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "13245c77",
   "metadata": {},
   "outputs": [],
   "source": [
    "clib_winograd_oneDNN = np.ctypeslib.load_library(\"libdnnl_winograd.so\", \".\")\n",
    "conv_winograd_oneDNN = conv_c_API(clib_winograd_oneDNN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "070fcc0d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- (N,IC,IH,IW,OC): (  8, 256,  56,  56, 256); GFLOP    :     3.440\n",
      "---    ms/run/batch:     3.938 ± 0.049        ; GFLOP/sec:   873.397\n"
     ]
    }
   ],
   "source": [
    "evaluate_GFLOPs(image, filtr, conv_winograd_oneDNN, loops=100, preloops=10);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd272882",
   "metadata": {},
   "source": [
    "- **我们自己编写的初赛程序偶尔能比 Intel oneDNN 快一些**。它是一个相对比较简单的，实现了 Winograd 算法的程序。我们实现的 Winograd 是 $F(6, 3)$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "4a99dce5",
   "metadata": {},
   "outputs": [],
   "source": [
    "clib_winograd_f6x3 = np.ctypeslib.load_library(\"libwinograd_f6x3.so\", \".\")\n",
    "conv_winograd_f6x3 = conv_c_API(clib_winograd_f6x3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "7ea2acd1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- (N,IC,IH,IW,OC): (  8, 256,  56,  56, 256); GFLOP    :     3.440\n",
      "---    ms/run/batch:     3.708 ± 0.167        ; GFLOP/sec:   927.607\n"
     ]
    }
   ],
   "source": [
    "evaluate_GFLOPs(image, filtr, conv_winograd_f6x3, loops=100, preloops=10);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9218c05b",
   "metadata": {},
   "source": [
    "- **PyTorch** 作为最流行的机器学习框架之一，也提供了非常简单直观的底层卷积函数。PyTorch 会依据实际情形选择算法，可能是 Winograd 也可能是 Direct。注意在效率测评过程中，必须要将反向传播过程关闭 (`torch.set_grad_enabled`)。我们可以同时用下述函数给出 CPU 与 GPU 效率。需要说明的是，对于 PyTorch-GPU 情形，**GPU 显卡内存与 DRAM 的通讯时间很大，不可忽略**，如果所有计算都在 GPU 中进行，则 GFLOP/sec 会高到离谱。因此下面的代码会自 CPU 读取输入图像。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "1b68d19a",
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "def conv_pytorch(device, filtr):\n",
    "    IC, OC = filtr.shape[:2]\n",
    "    torch_layer = torch.nn.Conv2d(IC, OC, 3, device=device)\n",
    "    torch_layer.load_state_dict({\"weight\": torch.tensor(filtr), \"bias\": torch.zeros(OC)})\n",
    "    def inner_func(image, filtr):\n",
    "        result = torch_layer(torch.tensor(image, device=device)).cpu()\n",
    "        return result\n",
    "    return inner_func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "5457d238",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- (N,IC,IH,IW,OC): (  8, 256,  56,  56, 256); GFLOP    :     3.440\n",
      "---    ms/run/batch:     2.677 ± 0.039        ; GFLOP/sec:  1284.805\n"
     ]
    }
   ],
   "source": [
    "evaluate_GFLOPs(image, filtr, conv_pytorch(\"cpu\", filtr), loops=100);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "e13062a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# execuate this code if gpu is available\n",
    "# evaluate_GFLOPs(image, filtr, conv_pytorch(\"cuda\", filtr), loops=10);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b07f482b",
   "metadata": {},
   "source": [
    "总地来说，我们可以看出，对于当前的问题而言，**任何一个实用的 C/C++ 程序实现，都比 NumPy 实现能快上 15 倍**。这才能完美地完成 *拯救生命大挑战* ！\n",
    "\n",
    "不过囿于一些我自己也不了解的技术问题，使用 NumPy/ctypes 调用 C 程序，有时无法完全达到纯 C 程序的效率。这对实际应用可能影响不大，但对测评结果会有影响。因此，我们不在这里作更详尽的测评。在第三篇文档中，我们会完全讨论 C 代码，而不使用 Python。\n",
    "\n",
    "我们都知道，若单论 NumPy 处理矩阵计算，其实效率与库函数实现没有多大差别。到底什么影响了卷积的实现效率呢？我们会在以后的文档中，对初赛工作作说明，从而可以一窥高效库函数的实现中，开发者可能会考虑到的问题。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
